#! /bin/bash

# Description:
#	- To monitor stuck nvidia processes on GPU worker nodes.
#	- Run on each GPU worker node 
# Author:
#	Ran Du <duran@ihep.ac.cn>	2019-Nov-01	v1.0
# Usage:
# 	./<script_name.sh>

# Detect stuck nvidia processes, and send alarms if there are any.

ADMIN="duran@ihep.ac.cn"

processes=`ps aux | grep nvidia | grep -v grep`

IFS=$'\n'
WARNING=0
stuck_process=""
for process in ${processes}
do
	process_name=`echo ${process} | gawk '{print $11}'`
	#echo ${process_name}

	match_str_1=`echo ${process_name} | sed -n '/nvidia-modeset/p'`
	#echo $match_str_1
	
	match_str_2=`echo ${process_name} | sed -n '/irq\/[0-9]*-nvidia/p'`
	#echo $match_str_2

	match_str_3=`echo ${process_name} | sed -n '/^\[nvidia\]$/p'`
	#echo $match_str_3
	
	if [[ -z ${match_str_1} && -z ${match_str_2}  && -z ${match_str_3} ]] 
	then
		stuck_process="${stuck_process} ${process_name}"
		WARNING=1
		echo "WARNING is enabled."
	fi
	
done

if [[ ${WARNING} == 1 ]]
then
	node="`hostname -s`"
	time_str="`date +'%Y-%m-%d %H:%M:%S'`"

	mail_subject="GPU STUCK PROCs WARNING: ${node} on ${time_str}"
	mail_context="Dear Admin, \n\tNVIDIA GPU stuck processes(${stuck_process}) are detected on node ${node}, ${time_str}.\n\tPlease check ASAP.\n\nYours sincerely,\nPomelo"
	mail_admin=${ADMIN}

	echo -e ${mail_context} | mail -s ${mail_subject} ${mail_admin} 	
fi
